#!/usr/bin/env python
'''
Given sam file from Tophat, this program 
selected the mapping in a range and output the 
mapped reads to a fa file.  
'''

__author__ = "Yangho chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

from TophatOutFormat import SamLineInfo
from search_a_large_sorted_file import extract_record
import os
import getopt
import sys

def sam2UniquelyMappedfasta( iFileN, oFileN ):
    '''
    This function extracts the reads uniquely mapped
    to the reference in the sam file and output in 
    fasta format.    
    '''
    iFile = open ( iFileN, "r" )
    oFile = open ( oFileN, "w" )
    for line in iFile:
        sam_line = SamLineInfo( line )
        # TODO, check if this read is uniquely mapped
        QSCORE_THRESHOLD = 4
        QSCORE_UNKNOWN = 255
        if ( int( sam_line.FLAG ) & 0x0100 ) == 0 :
            if( sam_line.MAPQ < QSCORE_THRESHOLD or \
                sam_line.MAPQ == QSCORE_UNKNOWN ):
                oFile.write( '>' + sam_line.QNAME + '\n' )
                oFile.write( sam_line.SEQ + '\n' )
    iFile.close()
    oFile.close()

def extract_mapped_reads( iFileN, oFileN, range ):
    tmpFileN = iFileN + str( range[0] ) + str( range[1] ) + ".tmp"
    mapped_pos_col_id = 3
    extract_record( iFileN, tmpFileN, mapped_pos_col_id , range, 0 )
    sam2UniquelyMappedfasta( tmpFileN, oFileN )
    os.remove( tmpFileN )

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print "usage: ", sys.argv[0]
    print "The four parameters should be "
    print "(1) the input sam file path"
    print "(2) the output .fa file path"
    print "(3) lower bound of the extracted alignments' left end"
    print "(4) upper bound of the extracted alignments' left"
    print str( len( sys.argv ) - 1 ) + " parameters have been provided."
    sys.exit( 2 )

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error:
        usage()
        sys.exit( 2 );
    # process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            usage()
            sys.exit();
    if len( args ) == 4:
        iFileN = args[0]
        oFileN = args[1]
        range = [int( args[2] ), int( args[3] )]
        if( os.path.exists( iFileN ) ):
            extract_mapped_reads( iFileN, oFileN, range )
        else:
            print iFileN, "doesn't exist."
    else:
        usage()

if __name__ == '__main__':
    main()
